import os
from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tqdm import tqdm

api_key = "6145f7bf-4c4c-4776-800e-fd20eee6ea59"

pc = Pinecone(api_key=api_key)

index_name = "mnist-index"
if index_name not in pc.list_indexes().names():
    pc.create_index(
        name=index_name,
        dimension=64,  
        metric='euclidean', 
        spec=ServerlessSpec(cloud='aws', region='us-east-1')  
    )

index = pc.Index(index_name)

digits = load_digits()
X, y = digits.data, digits.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

batch_size = 1000
for i in range(0, len(X_train), batch_size):
    batch_vectors = [{"id": str(i + j), "values": X_train[i + j].tolist(), "metadata": {"label": int(y_train[i + j])}}
                     for j in range(min(batch_size, len(X_train) - i))]
    index.upsert(batch_vectors)

y_pred = []
k = 11
print("测试过程进度：")
for x in tqdm(X_test, desc="正在测试"):
    try:
        query_response = index.query(vector=x.tolist(), top_k=k, include_metadata=True)
        neighbors = query_response['matches']

        if not neighbors:
            print("未找到任何邻居，查询向量:", x.tolist())
            y_pred.append(-1)  
            continue
        
        neighbor_labels = [match['metadata']['label'] for match in neighbors]
        
        predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
        y_pred.append(predicted_label)
    except Exception as e:
        print(f"查询时出错：{e}")
        y_pred.append(-1) 

y_pred_filtered = [pred for pred in y_pred if pred != -1]
y_test_filtered = [y_test[i] for i in range(len(y_pred)) if y_pred[i] != -1]

accuracy = accuracy_score(y_test_filtered, y_pred_filtered)
print(f"使用 Pinecone 和 KNN 模型的准确率为: {accuracy:.4f}")

